import math
import random
from collections import Counter

def euclidean_distance(x1, x2):
    if len(x1) != len(x2):
        raise ValueError("两个样本的特征维度必须一致")
    dist_sq = sum((a - b) ** 2 for a, b in zip(x1, x2))
    return math.sqrt(dist_sq)

def knn_predict(x_test, X_train, y_train, k=3, task="classification"):
    n_train = len(X_train)
    if n_train == 0:
        raise ValueError("训练集不能为空")
    if len(y_train) != n_train:
        raise ValueError("训练集特征与标签长度必须一致")
    if k <= 0 or k > n_train:
        raise ValueError(f"k必须满足 0 < k ≤ 训练样本数（当前训练样本数={n_train}）")
    if task not in ["classification", "regression"]:
        raise ValueError("task只能是'classification'或'regression'")

    distances = []
    for i in range(n_train):
        x_train = X_train[i]
        y_train_i = y_train[i]
        dist = euclidean_distance(x_test, x_train)
        distances.append((dist, y_train_i))

    distances.sort(key=lambda x: x[0])
    k_neighbors = distances[:k]

    k_labels = [neighbor[1] for neighbor in k_neighbors]
    if task == "classification":
        vote_result = Counter(k_labels).most_common(1)[0][0]
        return vote_result
    else:
        regression_result = sum(k_labels) / len(k_labels)
        return regression_result

def KNN(X_test, X_train, y_train, k=3, task="classification"):
    if len(X_test) == 0:
        raise ValueError("测试集不能为空")
    predictions = [knn_predict(x, X_train, y_train, k, task) for x in X_test]
    return predictions

if __name__ == "__main__":
    print("=" * 50)
    print("测试1：KNN分类任务（鸢尾花简化数据集）")
    print("=" * 50)
    X_train_classify = [
        [1.4, 0.2], [1.3, 0.2], [1.5, 0.2], [1.4, 0.3], [1.6, 0.2],
        [4.5, 1.5], [4.2, 1.3], [4.3, 1.3], [4.4, 1.2], [4.1, 1.1]
    ]
    y_train_classify = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
    X_test_classify = [[1.4, 0.25], [4.3, 1.4], [2.0, 0.5], [4.0, 1.0]]
    y_true_classify = [0, 1, 0, 1]

    k = 3
    predictions_classify = KNN(X_test_classify, X_train_classify, y_train_classify, k=k, task="classification")

    print(f"训练集：{len(X_train_classify)}个样本，k={k}")
    print(f"测试集样本：{X_test_classify}")
    print(f"真实标签：{y_true_classify}")
    print(f"预测标签：{predictions_classify}")
    correct = sum([1 for p, t in zip(predictions_classify, y_true_classify) if p == t])
    accuracy = correct / len(X_test_classify)
    print(f"分类准确率：{accuracy:.2f}")

    print("\n" + "=" * 50)
    print("测试2：KNN回归任务（房价预测简化数据集）")
    print("=" * 50)
    X_train_regress = [
        [50, 1], [60, 2], [70, 2], [80, 3], [90, 3],
        [100, 4], [110, 4], [120, 5], [130, 5], [140, 6]
    ]
    y_train_regress = [80, 95, 110, 125, 140, 155, 170, 185, 200, 215]
    X_test_regress = [[65, 2], [95, 3], [115, 4]]
    y_true_regress = [102.5, 147.5, 177.5]

    predictions_regress = KNN(X_test_regress, X_train_regress, y_train_regress, k=k, task="regression")

    print(f"训练集：{len(X_train_regress)}个样本，k={k}")
    print(f"测试集样本：{X_test_regress}")
    print(f"真实房价：{y_true_regress}")
    print(f"预测房价：{[round(p, 2) for p in predictions_regress]}")
    mse = sum([(p - t) ** 2 for p, t in zip(predictions_regress, y_true_regress)]) / len(X_test_regress)
    print(f"回归MSE（越小越好）：{mse:.2f}")